# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: xguo

import enum
import json
import logging
import time
import zlib

from tornado import gen
from tornado.concurrent import Future
from tornado.ioloop import IOLoop
from tornado.websocket import websocket_connect

from coin.base.timestamp import get_timestamp
from coin.exchange.base.kr_rest.auth_base import AuthKey
from coin.exchange.base.rest.task_scheduler import Scheduler
from coin.exchange.okex_v3.kr_rest.native_private_client import OkexAuth

OKEX_V3_WSS_URL = 'wss://real.okex.com:8443/ws/v3'


class OperationType(enum.Enum):
  SUBSCRIBE = 'subscribe'
  UNSUBSCRIBE = 'unsubscribe'
  LOGIN = 'login'
  PING = 'ping'


def parse_okex_raw_ws_msg(raw_msg):
  logger = logging.getLogger(__name__)
  dec_msg = raw_msg
  try:
    dec_msg = zlib.decompress(raw_msg, -zlib.MAX_WBITS)
  except zlib.error:
    logger.warning('Okex ws message is not compressed!')

  if dec_msg == b'pong':
    return {'event': 'ping', 'success': True}
  else:
    return json.loads(dec_msg)


class OkexWsPrivateClientBase(object):
  def __init__(
      self,
      *,
      key_file,
      ioloop=None,
      notification_callback=None,
      close_callback=None,
      on_ready=None,
      timeout_sec=10,
      api_url=None,
  ):
    self._api_url = api_url or OKEX_V3_WSS_URL
    self._ws = None

    self._key = AuthKey.from_file(key_file)
    self._auth = OkexAuth(self._key.key_file)
    self._key.passphrase = self._key.get_value('passphrase')

    self._ioloop = ioloop or IOLoop.current()
    self._close_callback = close_callback or (lambda: None)
    self._notification_callback = notification_callback or (lambda x: None)
    self._on_ready = on_ready or (lambda: None)

    self._pending_future = None
    self._subscribed_channels = set()
    self._scheduler = Scheduler()

    self._ready = False
    self._stopped = False

    self._last_received_message_timestamp = None
    self._ping_period_sec = 10
    self._timeout = timeout_sec * 1e9
    self._num_timeouts = 0
    self._num_invalid_msg = 0

    self._init_params = dict(
        key_file=key_file,
        ioloop=ioloop,
        notification_callback=notification_callback,
        close_callback=close_callback,
        on_ready=on_ready,
        timeout_sec=timeout_sec,
        api_url=api_url,
    )
    self._logger = logging.getLogger(__name__)

  @property
  def last_received_message_timestamp(self):
    return self._last_received_message_timestamp

  def is_ready(self):
    return self._ready

  def _get_sign(self, timestamp):
    str_list = [timestamp, 'GET', '/users/self/verify']
    str_to_be_signed = ''.join(str_list)
    sign = self._auth.get_sign(str_to_be_signed)
    return sign

  def start(self):
    self._ioloop.add_callback(self._connect)
    self._install_periodic_task()

  def _install_periodic_task(self):
    @gen.coroutine
    def send_ping():
      response = yield self.ping()
      self._logger.debug('Event pong! %s', response)

    def check_timeout():
      timestamp = get_timestamp()
      future = self._pending_future
      if future is None:
        return

      elapsed_time = timestamp - future.timestamp
      if elapsed_time > self._timeout:
        future.set_exception(
            TimeoutError('Websocket timeout error! %s, %s' %
                         (elapsed_time / 1e9, future.operation)))
        self._pending_future = None

    self._scheduler.add_periodic_task(self._ping_period_sec, send_ping, immediate_first_run=False)
    self._scheduler.add_periodic_task(1, check_timeout, immediate_first_run=False)

  @gen.coroutine
  def _connect(self):
    # noinspection PyBroadException
    try:
      self._logger.info('%s try to connect to the OkexV3 websocket.', type(self).__name__)
      self._ws = yield websocket_connect(self._api_url)
      self._logger.info('OkexV3 websocket connected!')
      self._ioloop.add_callback(self._run)

      self._logger.info('Send auth info!')
      response = yield self.login()
      if response == {'event': 'login', 'success': True}:
        self._logger.info('OkexV3 WS API login successfully!')
      else:
        self._logger.error('Failed to login!')
        self.stop()
        return

      self._ready = True
      self._on_ready()
      self._scheduler.start()
      self._logger.info('OkexWsPrivateClient is ready.')
    except Exception:
      self._logger.exception('%s connect to websocket failed!', type(self).__name__)
      self.stop()

  @gen.coroutine
  def _run(self):
    while not self._stopped:
      msg = None
      # noinspection PyBroadException
      try:
        msg = yield self._ws.read_message()
      except Exception:
        self._logger.exception('Error reading ws message')
        self.stop()

      if msg is None:
        self._logger.error('Websocket is closed!')
        self.stop()
        return

      if len(msg) == 0:
        self._logger.info('Empty message received!')
        continue

      # noinspection PyBroadException
      try:
        self.on_message(msg)
      except Exception:
        self._logger.exception('Exception raised when process message!')

  def on_message(self, raw_msg: bytes):
    self._last_received_message_timestamp = get_timestamp()
    if self._stopped:
      return

    # noinspection PyBroadException
    data = None
    try:
      data = parse_okex_raw_ws_msg(raw_msg)
      event = data.get('event')
      future = self._pending_future
      if event:
        if future:
          self._pending_future = None
          future.set_result(data)
        else:
          self._logger.error('Ignored message: %s', data)
      else:
        self._notification_callback(data)
    except Exception:
      self._logger.exception('Invalid json string!\n%s\n%s', raw_msg, data)
      self._num_invalid_msg += 1
      if self._num_invalid_msg > 10:
        self._logger.error('Too many invalid msg received! Restart private client')
        self.stop()
      return

  def send_request(self, operation: OperationType, args):
    if self._pending_future:
      self._logger.error('There are futures that is not resolved yet! %s\n%s,%s',
                         self._pending_future.operation,
                         operation,
                         args)
      return

    message = {
        'op': operation.value,
        'args': args,
    }

    future = Future()
    future.operation = [operation, args]
    future.timestamp = get_timestamp()
    self._ws.write_message(json.dumps(message))
    self._pending_future = future
    return future

  def login(self):
    operation = OperationType.LOGIN
    timestamp = '%.03f' % time.time()
    args = [self._key.access_key, self._key.passphrase, timestamp, self._get_sign(timestamp)]
    fut = self.send_request(operation, args)
    return fut

  def ping(self):
    if self._ws:
      self._ws.write_message('ping')
      future = Future()
      future.operation = OperationType.PING
      future.timestamp = get_timestamp()
      self._pending_future = future
      return future

  def stop(self):
    if self._stopped:
      return
    self._ready = False
    self._ioloop.add_callback(self._async_stop)
    self._logger.info('Okex ws private client (%s) stopped!', id(self))

  @gen.coroutine
  def _async_stop(self):
    while self._pending_future:
      self._logger.info('Wait until pending future is resolved!')
      yield gen.sleep(1)

    if self._ws:
      for channel in self._subscribed_channels:
        response = yield self.send_request(OperationType.UNSUBSCRIBE, [channel])
        self._logger.info('Channel unsubscribed!\n%s', response)

    self._stopped = True
    if self._ws is not None:
      self._ws.close()
      self._ws = None
    self._close_callback()

  def clone(self):
    return type(self)(**self._init_params)


class OkexWsSpotPrivateClient(OkexWsPrivateClientBase):
  def subscribe_spot_balance(self, currency):
    args = ['spot/account:' + currency]
    return self.send_request(OperationType.SUBSCRIBE, args)

  def subscribe_spot_order(self, product):
    args = ['spot/order:' + product]
    return self.send_request(OperationType.SUBSCRIBE, args)

  @gen.coroutine
  def prepare_spot_subscription(self, products):
    for product in products:
      response = yield self.subscribe_spot_order(product)
      event = response['event']
      if event != 'error':
        self._logger.info('Subscribe order response: %s\n', response)
        channel = response.get('channel')
        self._subscribed_channels.add(channel)
      else:
        self._logger.info('Balance subscription failed: %s', response)
        self.stop()
        return
